#include "Socket.hpp"
#include "Epoll.hpp"
#include "Log.hpp"
#include <string>
#include <unordered_map>
#include <atomic>
#include <iostream>
#include <cstring>
#include <sstream>
#include <netdb.h>
#include <fcntl.h>
#include <memory>

// 读写数据的缓冲区大小
const int BUFFER_SIZE = 8192;

// 连接类型枚举
enum ConnectionType
{
    CLIENT, // 客户端连接
    SERVER  // 到目标服务器的连接
};

// 连接状态枚举
enum ConnectionState
{
    READING_REQUEST,   // 读取客户端请求
    CONNECTING_SERVER, // 连接到目标服务器
    SENDING_REQUEST,   // 发送请求到目标服务器
    READING_RESPONSE,  // 读取目标服务器响应
    SENDING_RESPONSE,  // 发送响应到客户端
    CLOSED             // 连接已关闭
};

// HTTP连接的结构体
struct Connection
{
    int fd;                  // 文件描述符
    ConnectionType type;     // 连接类型
    ConnectionState state;   // 连接状态
    std::string inBuffer;    // 输入缓冲区
    std::string outBuffer;   // 输出缓冲区
    int pairFd;              // 配对连接的文件描述符
    std::string host;        // 目标主机名
    int port;                // 目标端口
    std::string httpVersion; // HTTP版本

    Connection(int _fd, ConnectionType _type)
        : fd(_fd), type(_type), state(_type == CLIENT ? READING_REQUEST : CONNECTING_SERVER),
          pairFd(-1), port(80) {}

    ~Connection()
    {
        close(fd);
    }
};

class HttpProxy
{
public:
    HttpProxy(int port) : proxyPort(port), running(false), maxEvents(128), listenFd(-1)
    {
    }

    ~HttpProxy()
    {
        running = false;
        for (auto &pair : connections)
        {
            if (pair.second)
            {
                close(pair.first);
                delete pair.second;
            }
        }
        connections.clear();
        if (events)
        {
            delete[] events;
            events = nullptr;
        }
        logMessage(NORMAL, "HTTP代理服务器已停止");
    }

    // 启动代理服务器
    bool start()
    {
        try
        {
            listenFd = sock.Socket();
            sock.Bind(listenFd, proxyPort);
            sock.Listen(listenFd);
            epoll.Create();
            epoll.Control(listenFd, EPOLL_CTL_ADD, EPOLLIN);
            events = new epoll_event[maxEvents];
            running = true;
            logMessage(NORMAL, "HTTP代理服务器已启动，监听端口: %d", proxyPort);
            eventLoop();
            return true;
        }
        catch (const std::exception &e)
        {
            logMessage(ERROR, "启动HTTP代理服务器失败: %s", e.what());
            return false;
        }
    }

private:
    // 设置文件描述符为非阻塞模式
    void setNonBlocking(int fd)
    {
        int flags = fcntl(fd, F_GETFL, 0);
        if (flags < 0)
        {
            logMessage(ERROR, "获取文件描述符标志失败");
            return;
        }
        if (fcntl(fd, F_SETFL, flags | O_NONBLOCK) < 0)
        {
            logMessage(ERROR, "设置非阻塞模式失败");
        }
    }
    // 主循环处理到来事件
    void eventLoop()
    {
        while (running)
        {
            int numEvents = epoll.Wait(maxEvents, events);
            for (int i = 0; i < numEvents; i++)
            {
                int fd = events[i].data.fd;
                if (fd == listenFd)
                {
                    handleAccept();
                    continue;
                }
                auto it = connections.find(fd);
                if (it == connections.end())
                {
                    epoll.Control(fd, EPOLL_CTL_DEL, 0);
                    close(fd);
                    continue;
                }
                Connection *conn = it->second;
                if (events[i].events & (EPOLLERR | EPOLLHUP))
                {
                    closeConnection(conn);
                    continue;
                }
                if (events[i].events & EPOLLIN)
                {
                    if (conn->type == CLIENT)
                    {
                        switch (conn->state)
                        {
                        case READING_REQUEST:
                            if (handleClientRead(conn))
                            {
                                parseRequest(conn);
                            }
                            break;
                        default:
                            break;
                        }
                    }
                    else
                    {
                        switch (conn->state)
                        {
                        case READING_RESPONSE:
                            handleServerRead(conn);
                            break;
                        default:
                            break;
                        }
                    }
                }
                if (events[i].events & EPOLLOUT)
                {
                    if (conn->type == CLIENT)
                    {
                        switch (conn->state)
                        {
                        case SENDING_RESPONSE:
                            handleClientWrite(conn);
                            break;
                        default:
                            break;
                        }
                    }
                    else
                    {
                        switch (conn->state)
                        {
                        case CONNECTING_SERVER:
                            handleServerConnect(conn);
                            break;
                        case SENDING_REQUEST:
                            handleServerWrite(conn);
                            break;
                        default:
                            break;
                        }
                    }
                }
            }
        }
    }

    // 处理接受新的客户端连接
    void handleAccept()
    {
        int err = 0;
        int clientFd = sock.Accept(listenFd, err);
        if (clientFd > 0)
        {
            setNonBlocking(clientFd);
            Connection *conn = new Connection(clientFd, CLIENT);
            connections[clientFd] = conn;
            epoll.Control(clientFd, EPOLL_CTL_ADD, EPOLLIN | EPOLLET);
            logMessage(NORMAL, "接受新客户端连接: [%d]", clientFd);
        }
        else if (err != EAGAIN && err != EWOULDBLOCK && err != EINTR)
        {
            logMessage(ERROR, "接受连接出错: %s", strerror(err));
        }
    }

    // 处理客户端连接的读取
    bool handleClientRead(Connection *conn)
    {
        char buffer[BUFFER_SIZE];
        bool requestComplete = false;
        while (true)
        {
            int bytesRead = recv(conn->fd, buffer, BUFFER_SIZE - 1, 0);
            if (bytesRead > 0)
            {
                buffer[bytesRead] = '\0';
                conn->inBuffer.append(buffer, bytesRead);
                size_t headerEnd = conn->inBuffer.find("\r\n\r\n");
                if (headerEnd != std::string::npos)
                {
                    requestComplete = true;
                }
            }
            else if (bytesRead == 0)
            {
                logMessage(NORMAL, "客户端关闭连接: [%d]", conn->fd);
                closeConnection(conn);
                return false;
            }
            else
            {
                if (errno == EAGAIN || errno == EWOULDBLOCK)
                {
                    break;
                }
                else if (errno == EINTR)
                {
                    continue;
                }
                else
                {
                    logMessage(ERROR, "从客户端读取数据出错: %s", strerror(errno));
                    closeConnection(conn);
                    return false;
                }
            }
        }
        return requestComplete;
    }

    // 解析HTTP请求并连接到目标服务器
    void parseRequest(Connection *conn)
    {
        size_t headerEnd = conn->inBuffer.find("\r\n\r\n");
        if (headerEnd == std::string::npos)
        {
            return;
        }
        std::string requestData = conn->inBuffer.substr(0, headerEnd + 4);
        std::istringstream requestStream(requestData);
        std::string method, url;
        requestStream >> method >> url >> conn->httpVersion;
        if (method != "GET" && method != "POST" && method != "HEAD")
        {
            std::string errorResponse = "HTTP/1.1 501 Not Implemented\r\nContent-Length: 0\r\n\r\n";
            conn->outBuffer = errorResponse;
            conn->state = SENDING_RESPONSE;
            epoll.Control(conn->fd, EPOLL_CTL_MOD, EPOLLIN | EPOLLOUT | EPOLLET);
            logMessage(ERROR, "不支持的方法: %s", method.c_str());
            return;
        }
        size_t hostPos = requestData.find("Host: ");
        if (hostPos != std::string::npos)
        {
            size_t hostEndPos = requestData.find("\r\n", hostPos);
            if (hostEndPos != std::string::npos)
            {
                conn->host = requestData.substr(hostPos + 6, hostEndPos - hostPos - 6);
                size_t colonPos = conn->host.find(":");
                if (colonPos != std::string::npos)
                {
                    conn->port = std::stoi(conn->host.substr(colonPos + 1));
                    conn->host = conn->host.substr(0, colonPos);
                }
            }
        }
        if (url.find("http://") == 0)
        {
            std::string fullUrl = url;
            url = fullUrl.substr(7);
            size_t pathPos = url.find("/");
            if (pathPos != std::string::npos)
            {
                std::string hostPort = url.substr(0, pathPos);
                url = url.substr(pathPos);
                size_t colonPos = hostPort.find(":");
                if (colonPos != std::string::npos)
                {
                    conn->host = hostPort.substr(0, colonPos);
                    conn->port = std::stoi(hostPort.substr(colonPos + 1));
                }
                else
                {
                    conn->host = hostPort;
                }
            }
        }
        if (conn->host.empty())
        {
            std::string errorResponse = "HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\n\r\n";
            conn->outBuffer = errorResponse;
            conn->state = SENDING_RESPONSE;
            epoll.Control(conn->fd, EPOLL_CTL_MOD, EPOLLIN | EPOLLOUT | EPOLLET);
            logMessage(ERROR, "请求中没有Host字段或URL中没有主机信息");
            return;
        }
        std::string modifiedRequest = method + " " + url + " " + conn->httpVersion + "\r\n";
        modifiedRequest += requestData.substr(requestData.find("\r\n") + 2);
        connectToServer(conn, modifiedRequest);
    }

    // 连接到目标服务器
    void connectToServer(Connection *clientConn, const std::string &request)
    {
        int serverFd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0);
        if (serverFd < 0)
        {
            logMessage(ERROR, "创建到目标服务器的socket失败");
            closeConnection(clientConn);
            return;
        }
        struct hostent *server = gethostbyname(clientConn->host.c_str());
        if (server == nullptr)
        {
            logMessage(ERROR, "解析主机名失败: %s", clientConn->host.c_str());
            close(serverFd);
            closeConnection(clientConn);
            return;
        }
        sockaddr_in serverAddr;
        memset(&serverAddr, 0, sizeof(serverAddr));
        serverAddr.sin_family = AF_INET;
        memcpy(&serverAddr.sin_addr.s_addr, server->h_addr, server->h_length);
        serverAddr.sin_port = htons(clientConn->port);
        Connection *serverConn = new Connection(serverFd, SERVER);
        serverConn->outBuffer = request;
        serverConn->pairFd = clientConn->fd;
        clientConn->pairFd = serverFd;
        connections[serverFd] = serverConn;
        int result = connect(serverFd, (struct sockaddr *)&serverAddr, sizeof(serverAddr));
        if (result < 0)
        {
            if (errno == EINPROGRESS)
            {
                epoll.Control(serverFd, EPOLL_CTL_ADD, EPOLLOUT | EPOLLET);
                logMessage(NORMAL, "正在连接到目标服务器: %s:%d", clientConn->host.c_str(), clientConn->port);
            }
            else
            {
                logMessage(ERROR, "连接到目标服务器失败: %s", strerror(errno));
                closeConnection(serverConn);
                closeConnection(clientConn);
            }
        }
        else
        {
            serverConn->state = SENDING_REQUEST;
            epoll.Control(serverFd, EPOLL_CTL_ADD, EPOLLOUT | EPOLLET);
        }
    }

    // 处理服务器连接建立
    void handleServerConnect(Connection *conn)
    {
        int error = 0;
        socklen_t len = sizeof(error);
        if (getsockopt(conn->fd, SOL_SOCKET, SO_ERROR, &error, &len) < 0 || error != 0)
        {
            logMessage(ERROR, "连接到目标服务器失败: %s", strerror(error));
            if (conn->pairFd != -1)
            {
                auto it = connections.find(conn->pairFd);
                if (it != connections.end())
                {
                    Connection *clientConn = it->second;
                    std::string errorResponse = "HTTP/1.1 502 Bad Gateway\r\nContent-Length: 0\r\n\r\n";
                    clientConn->outBuffer = errorResponse;
                    clientConn->state = SENDING_RESPONSE;
                    epoll.Control(clientConn->fd, EPOLL_CTL_MOD, EPOLLIN | EPOLLOUT | EPOLLET);
                }
            }
            closeConnection(conn);
            return;
        }
        conn->state = SENDING_REQUEST;
        epoll.Control(conn->fd, EPOLL_CTL_MOD, EPOLLOUT | EPOLLET);
        logMessage(NORMAL, "成功连接到目标服务器，准备发送请求 [%d]", conn->fd);
    }

    // 处理向服务器写入请求数据
    void handleServerWrite(Connection *conn)
    {
        while (!conn->outBuffer.empty())
        {
            int bytesSent = send(conn->fd, conn->outBuffer.c_str(), conn->outBuffer.size(), 0);
            if (bytesSent > 0)
            {
                conn->outBuffer.erase(0, bytesSent);
            }
            else if (bytesSent < 0)
            {
                if (errno == EAGAIN || errno == EWOULDBLOCK)
                {
                    break;
                }
                else if (errno == EINTR)
                {
                    continue;
                }
                else
                {
                    logMessage(ERROR, "向目标服务器发送请求失败: %s", strerror(errno));
                    closeConnection(conn);
                    return;
                }
            }
        }
        if (conn->outBuffer.empty())
        {
            conn->state = READING_RESPONSE;
            epoll.Control(conn->fd, EPOLL_CTL_MOD, EPOLLIN | EPOLLET);
            logMessage(NORMAL, "请求已发送给目标服务器，等待响应 [%d]", conn->fd);
        }
    }

    // 处理从服务器读取响应数据
    void handleServerRead(Connection *conn)
    {
        char buffer[BUFFER_SIZE];
        while (true)
        {
            int bytesRead = recv(conn->fd, buffer, BUFFER_SIZE - 1, 0);
            if (bytesRead > 0)
            {
                buffer[bytesRead] = '\0';
                auto it = connections.find(conn->pairFd);
                if (it != connections.end())
                {
                    Connection *clientConn = it->second;
                    clientConn->outBuffer.append(buffer, bytesRead);
                    clientConn->state = SENDING_RESPONSE;
                    epoll.Control(clientConn->fd, EPOLL_CTL_MOD, EPOLLIN | EPOLLOUT | EPOLLET);
                }
            }
            else if (bytesRead == 0)
            {
                logMessage(NORMAL, "目标服务器关闭连接，响应结束 [%d]", conn->fd);
                auto it = connections.find(conn->pairFd);
                if (it != connections.end())
                {
                    Connection *clientConn = it->second;
                    if (!clientConn->outBuffer.empty())
                    {
                        clientConn->pairFd = -1;
                    }
                }
                closeConnection(conn);
                return;
            }
            else
            {
                if (errno == EAGAIN || errno == EWOULDBLOCK)
                {
                    break;
                }
                else if (errno == EINTR)
                {
                    continue;
                }
                else
                {
                    logMessage(ERROR, "从目标服务器读取响应失败: %s", strerror(errno));
                    closeConnection(conn);
                    return;
                }
            }
        }
    }

    // 处理向客户端写入响应数据
    void handleClientWrite(Connection *conn)
    {
        while (!conn->outBuffer.empty())
        {
            int bytesSent = send(conn->fd, conn->outBuffer.c_str(), conn->outBuffer.size(), 0);
            if (bytesSent > 0)
            {
                conn->outBuffer.erase(0, bytesSent);
            }
            else if (bytesSent < 0)
            {
                if (errno == EAGAIN || errno == EWOULDBLOCK)
                {
                    break;
                }
                else if (errno == EINTR)
                {
                    continue;
                }
                else
                {
                    logMessage(ERROR, "向客户端发送响应失败: %s", strerror(errno));
                    closeConnection(conn);
                    return;
                }
            }
        }
        if (conn->outBuffer.empty())
        {
            if (conn->pairFd == -1)
            {
                closeConnection(conn);
            }
            else
            {
                conn->state = READING_REQUEST;
                conn->inBuffer.clear();
                epoll.Control(conn->fd, EPOLL_CTL_MOD, EPOLLIN | EPOLLET);
                logMessage(NORMAL, "响应已发送给客户端，等待下一个请求 [%d]", conn->fd);
            }
        }
    }

    // 关闭连接
    void closeConnection(Connection *conn)
    {
        if (!conn)
            return;
        int fd = conn->fd;
        epoll.Control(fd, EPOLL_CTL_DEL, 0);
        if (conn->pairFd != -1)
        {
            auto it = connections.find(conn->pairFd);
            if (it != connections.end())
            {
                it->second->pairFd = -1;
            }
        }
        connections.erase(fd);
        delete conn;
        logMessage(NORMAL, "关闭连接 [%d]", fd);
    }

    // 代理服务器监听的端口
    int proxyPort;

    // Socket和Epoll封装对象
    Sock sock;
    Epoll epoll;

    // 控制主循环的标志
    std::atomic<bool> running;

    // Epoll事件数组
    struct epoll_event *events;
    int maxEvents;

    // 连接管理映射
    std::unordered_map<int, Connection *> connections;

    int listenFd;
};
